/*
 * @Author: Wangjun
 * @Date: 2021-05-15 22:43:20
 * @LastEditTime: 2022-08-23 10:33:01
 * @LastEditors: Wangjun
 * @Description:
 * @FilePath: \fan_wind_powerd:\go\src\gitee.com\haodreams\golib\wrap\wrap.go
 * hnxr
 */
package wrap

import (
	"sync"
	"sync/atomic"
	"time"
)

// WaitGroupWrapper .
type WaitGroupWrapper struct {
	sync.WaitGroup
	num int32
	max int32 //最大线程数
}

// NewWaitGroupWrapper 支持设置最大线程数的wg
func NewWaitGroupWrapper(maxThread int) *WaitGroupWrapper {
	wg := new(WaitGroupWrapper)
	wg.max = int32(maxThread)
	return wg
}

// Wrap cb 回调函数
// 回调函数需要注意的是 回调函数内部的变量是引用变量，需要注意创建新变量复制
func (m *WaitGroupWrapper) Wrap(cb func()) {
	if m.max > 0 {
		atomic.AddInt32(&m.num, 1)
	}

	m.Add(1)
	go func() {
		defer m.Done()
		defer func() {
			if m.max > 0 {
				atomic.AddInt32(&m.num, -1)
			}
		}()
		cb()
	}()

	if m.max > 0 {
		for atomic.LoadInt32(&m.num) > m.max {
			time.Sleep(time.Millisecond * 500)
		}
	}
}
